import streamlit as st
import os
from PIL import Image
import pandas as pd

from models.vllm_model import VLLMModel
from utils.helpers import save_upload_file

def show_conversation():
    """显示对话模式页面"""
    
    st.title("对话模式")
    
    # 初始化会话状态
    if "vllm_model" not in st.session_state:
        st.session_state.vllm_model = None
    
    # 初始化对话历史
    if "conversation_history" not in st.session_state:
        st.session_state.conversation_history = []
    
    # 初始化当前图片
    if "current_image_path" not in st.session_state:
        st.session_state.current_image_path = None
    
    # 侧边栏配置
    st.sidebar.header("模型配置")
    
    # Ollama服务地址配置
    ollama_url = st.sidebar.text_input("Ollama服务地址", value="http://localhost:11434")
    
    # 模型选择
    try:
        if st.session_state.vllm_model is None:
            # 初始化VLLM模型
            st.session_state.vllm_model = VLLMModel(base_url=ollama_url)
        
        # 刷新模型列表
        if st.sidebar.button("刷新模型列表"):
            refresh_placeholder = st.sidebar.empty()
            refresh_placeholder.info("正在获取可用模型...")
            try:
                available_models = st.session_state.vllm_model.get_available_models()
                st.session_state.available_models = available_models
                refresh_placeholder.success("模型列表已更新")
            except Exception as e:
                refresh_placeholder.error("获取模型列表失败")
    except Exception as e:
        st.sidebar.error(f"连接Ollama服务失败: {str(e)}")
        st.sidebar.warning("请确保Ollama服务已启动，并检查服务地址是否正确")
        available_models = ["qwen2.5-vl"]
    
    # 如果会话状态中有模型列表，则使用它
    if "available_models" in st.session_state:
        available_models = st.session_state.available_models
    else:
        available_models = ["qwen2.5-vl"]
    
    # 选择模型
    selected_model = st.sidebar.selectbox("选择VLLM模型", available_models)
    
    # 更新模型配置
    if st.sidebar.button("应用配置"):
        try:
            st.session_state.vllm_model = VLLMModel(base_url=ollama_url, model_name=selected_model)
            st.sidebar.success(f"成功配置模型: {selected_model}")
        except Exception as e:
            st.sidebar.error(f"配置模型失败: {str(e)}")
    
    # 清空对话历史按钮
    if st.sidebar.button("清空对话历史"):
        st.session_state.conversation_history = []
        st.session_state.current_image_path = None
    
    # 主界面
    st.header("产品缺陷对话")
    
    # 检查知识库和CLIP模型
    if "knowledge_base" not in st.session_state or st.session_state.knowledge_base is None:
        st.warning("请先在'知识库管理'页面加载CLIP模型和知识库")
        return
    
    # 获取可用的知识库
    kb_list = st.session_state.knowledge_base.get_product_list()
    
    if not kb_list:
        st.warning("未找到任何知识库，请先创建知识库")
        return
    
    # 选择产品
    selected_product = st.selectbox("选择产品", kb_list)
    
    # 创建两个选项卡：图片对话和文本问答
    tab1, tab2 = st.tabs(["📷 图片对话", "💬 文本问答"])
    
    # 图片对话选项卡
    with tab1:
        st.subheader("上传图片进行对话")
        # 上传图片
        uploaded_file = st.file_uploader("上传图片", type=["png", "jpg", "jpeg"])
        
        # 如果上传了新图片
        if uploaded_file is not None:
            # 显示上传的图片
            image = Image.open(uploaded_file)
            st.image(image, caption="当前图片", use_column_width=True)
            
            # 保存上传的图片
            upload_dir = os.path.join("./uploads", "conversation")
            os.makedirs(upload_dir, exist_ok=True)
            img_path = save_upload_file(uploaded_file, upload_dir)
            
            # 更新当前图片路径
            st.session_state.current_image_path = img_path
            
            # 设置对话模式为图片对话
            st.session_state.conversation_mode = "image"
    
    # 文本问答选项卡
    with tab2:
        st.subheader("文本问答")
        st.info(f"当前产品: {selected_product}")
        
        # 设置对话模式为文本问答
        if st.button("使用文本问答模式"):
            st.session_state.conversation_mode = "text"
            st.session_state.current_image_path = None
            st.success("已切换到文本问答模式")
            
    # 初始化对话模式
    if "conversation_mode" not in st.session_state:
        st.session_state.conversation_mode = "image"
    
    # 显示对话区域
    st.header("对话区域")
    
    # 显示对话历史
    for message in st.session_state.conversation_history:
        role = message["role"]
        content = message["content"]
        
        if role == "user":
            st.markdown(f"**👤 您**: {content}")
        else:
            st.markdown(f"**🤖 AI**: {content}")
    
    # 用户输入
    user_input = st.text_input("输入您的问题", key="user_input")
    
    # 当用户提交问题时
    if st.button("发送") and user_input:
        # 将用户问题添加到对话历史
        st.session_state.conversation_history.append({
            "role": "user",
            "content": user_input
        })
        
        # 加载知识库
        with st.spinner("正在处理..."):
            try:
                # 加载知识库
                st.session_state.knowledge_base.load_or_create(selected_product)
                
                # 检查对话模式
                if st.session_state.conversation_mode == "image":
                    # 检查是否有上传的图片
                    if st.session_state.current_image_path is None:
                        st.error("请先上传图片")
                        # 从对话历史中移除最后一条
                        st.session_state.conversation_history.pop()
                        return
                    
                    # 从知识库检索相似图片
                    retrieval_results = st.session_state.knowledge_base.search(st.session_state.current_image_path, top_k=5)
                    
                    # 调用VLLM模型进行对话（基于图片）
                    result = st.session_state.vllm_model.chat(
                        image_path=st.session_state.current_image_path,
                        user_query=user_input,
                        product_name=selected_product,
                        knowledge_results=retrieval_results
                    )
                else:  # 文本问答模式
                    # 使用用户问题进行文本检索
                    retrieval_results = st.session_state.knowledge_base.text_search(user_input, selected_product, top_k=5)
                    
                    # 调用VLLM模型进行对话（基于文本）
                    result = st.session_state.vllm_model.text_chat(
                        user_query=user_input,
                        product_name=selected_product,
                        knowledge_results=retrieval_results
                    )
                
                # 显示检索结果
                st.subheader("知识库检索结果")
                
                # 创建结果表格
                retrieval_df = []
                
                # 添加检索结果
                for i, res in enumerate(retrieval_results):
                    retrieval_df.append({
                        "序号": i + 1,
                        "图片路径": res.get("path", ""),
                        "标签": res.get("label", "未知"),
                        "相似度": f"{res.get('similarity', 0):.4f}"
                    })
                
                # 显示表格
                if retrieval_df:
                    df = pd.DataFrame(retrieval_df)
                    st.dataframe(df, use_container_width=True)
                    
                    # 显示检索到的图片
                    st.subheader("检索到的相似图片")
                    cols = st.columns(min(5, len(retrieval_results)))
                    for i, res in enumerate(retrieval_results[:5]):  # 最多显示5张图片
                        try:
                            img_path = res.get("path", "")
                            if os.path.exists(img_path):
                                with cols[i]:
                                    img = Image.open(img_path)
                                    st.image(img, caption=f"{res.get('label', '未知')} ({res.get('similarity', 0):.4f})", use_column_width=True)
                        except Exception as e:
                            st.error(f"无法加载图片 {img_path}: {str(e)}")
                else:
                    st.warning("未找到相似图片")
                
                # 获取并保存prompt
                retrieval_context = st.session_state.vllm_model._format_retrieval_results(retrieval_results)
                
                if st.session_state.conversation_mode == "image":
                    prompt = f"""你是一个产品表面缺陷检测专家，请根据用户的问题，分析图片中的{selected_product}产品。

以下是从相似图片知识库中检索到的结果供你参考:
{retrieval_context}

用户问题: {user_input}

请用中文回答用户的问题，回答要简洁专业。"""
                else:
                    prompt = f"""你是一个产品表面缺陷检测专家，请回答用户关于{selected_product}产品的问题。

以下是从知识库中检索到的相关图片结果供你参考:
{retrieval_context}

用户问题: {user_input}

请用中文回答用户的问题，回答要简洁专业。如果知识库中有相关的缺陷图片，请在回答中描述这些缺陷的特征。"""
                
                # 显示传给大模型的prompt
                with st.expander("查看传给大模型的Prompt"):
                    st.code(prompt, language="markdown")
                
                if result.get("success", False):
                    response = result.get("response", "抱歉，无法处理您的请求。")
                    
                    # 将AI回复添加到对话历史
                    st.session_state.conversation_history.append({
                        "role": "assistant",
                        "content": response
                    })
                    
                    # 刷新页面显示最新对话
                    st.rerun()
                else:
                    st.error(f"处理失败: {result.get('error', '未知错误')}")
                    
            except Exception as e:
                st.error(f"对话过程发生错误: {str(e)}")
    
    st.markdown("---")
    
    st.markdown("""
    ### 使用说明
    
    1. 选择已创建的产品知识库
    2. 上传产品图片
    3. 在对话区域输入问题，例如:
       - 这个产品有什么缺陷？
       - 这是什么类型的缺陷？
       - 这个产品可能存在哪些类型的缺陷？
       - 如何预防这种缺陷？
    
    系统会基于图像内容和知识库中的相似案例，给出专业的回答。
    """) 